package org.apache.dolphinscheduler.api.controller;

import static org.apache.dolphinscheduler.api.enums.Status.CREATE_MODEL_FILE_ERROR;
import static org.apache.dolphinscheduler.api.enums.Status.CREATE_MODEL_LIST_ERROR;
import static org.apache.dolphinscheduler.api.enums.Status.DELETE_MODEL_FILE_ERROR;
import static org.apache.dolphinscheduler.api.enums.Status.MODEL_QUERY_LIST_ERROR;
import static org.apache.dolphinscheduler.api.enums.Status.PUBLISH_MODEL_ERROR;
import static org.apache.dolphinscheduler.api.enums.Status.RESOURCE_NOT_EXIST;
import static org.apache.dolphinscheduler.api.enums.Status.UNPUBLISH_MODEL_ERROR;
import static org.apache.dolphinscheduler.api.enums.Status.UPLOAD_MODEL_FILE_ERROR;

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;

import org.apache.commons.io.IOUtils;
import org.apache.dolphinscheduler.api.aspect.AccessLogAnnotation;
import org.apache.dolphinscheduler.api.aspect.PermissionCheck;
import org.apache.dolphinscheduler.api.dto.model.ModelDto;
import org.apache.dolphinscheduler.api.dto.model.ModelFileDto;
import org.apache.dolphinscheduler.api.dto.model.PublishModelDto;
import org.apache.dolphinscheduler.api.exceptions.ApiException;
import org.apache.dolphinscheduler.api.service.ModelFileService;
import org.apache.dolphinscheduler.api.service.ModelService;
import org.apache.dolphinscheduler.api.utils.Result;
import org.apache.dolphinscheduler.common.constants.Constants;
import org.apache.dolphinscheduler.dao.entity.Model;
import org.apache.dolphinscheduler.dao.entity.ModelProfile;
import org.apache.dolphinscheduler.dao.entity.ModelVersion;
import org.apache.dolphinscheduler.dao.entity.User;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.web.bind.annotation.DeleteMapping;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.PutMapping;
import org.springframework.web.bind.annotation.RequestAttribute;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.ResponseStatus;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.Parameters;
import io.swagger.v3.oas.annotations.media.Schema;
import io.swagger.v3.oas.annotations.tags.Tag;

@Tag(name = "MODEL_TAG")
@RestController
@RequestMapping("model")
public class ModelController extends BaseController {

	@Autowired
	private ModelService modelService;
	
	@Autowired
	private ModelFileService modelFileService;
	
	@Operation(summary = "queryModelList", description = "QUERY_MODEL_LIST_NOTES")
    @GetMapping("/list")
    @ApiException(MODEL_QUERY_LIST_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:list,empower:dataset:list")
	public Result<Object> listModel(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, ModelDto dto){
		return modelService.queryModelByPaging(loginUser, dto);
	}
	
	@Operation(summary = "createModel", description = "CREATE_MODEL_NOTES")
    @PostMapping("/create")
    @ApiException(CREATE_MODEL_LIST_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:create,empower:dataset:create")
	public Result<Object> createModel(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, @RequestBody Model model){
		return modelService.createModel(loginUser, model);
	}
	
	@Operation(summary = "updateModel", description = "UPDATE_MODEL_NOTES")
    @PutMapping("/{id}")
    @ApiException(CREATE_MODEL_LIST_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:update,empower:dataset:update")
	public Result<Object> updateModel(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, 
									  @PathVariable(value = "id") Integer modelId,
									  @RequestBody Model model){
		return modelService.updateModel(loginUser, modelId, model);
	}
	
	@Operation(summary = "createProfile", description = "CREATE_PROFILE_NOTES")
    @PostMapping("/{id}/profile")
    @ApiException(CREATE_MODEL_LIST_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:profile:create,empower:dataset:profile:create")
	public Result<Object> createProfile(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, 
									@PathVariable(value = "id") Integer modelId, 
									@RequestBody ModelProfile profile){
		return modelService.createProfile(loginUser, modelId, profile);
	}
	
	@Operation(summary = "getProfile", description = "GET_PROFILE_NOTES")
    @GetMapping("/{id}/profile")
    @ApiException(CREATE_MODEL_LIST_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:profile:get,empower:dataset:profile:get")
	public Result<Object> getProfile(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, 
									@PathVariable(value = "id") Integer modelId){
		return modelService.getProfile(loginUser, modelId);
	}
	
	@Operation(summary = "createVersion", description = "CREATE_VERSION_NOTES")
    @PostMapping("/{id}/version")
    @ApiException(CREATE_MODEL_LIST_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:version:create,empower:dataset:version:create")
	public Result<Object> createVersion(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, 
									@PathVariable(value = "id") Integer modelId, 
									@RequestBody ModelVersion version){
		return modelService.createVersion(loginUser, modelId, version);
	}
	
	@Operation(summary = "listVersion", description = "LIST_VERSION_NOTES")
    @GetMapping("/{id}/version")
    @ApiException(CREATE_MODEL_LIST_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:version:list,empower:dataset:version:list")
	public Result<Object> listVersion(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, 
									@PathVariable(value = "id") Integer modelId){
		return modelService.listVersion(loginUser, modelId);
	}
	
	@Operation(summary = "updateProfile", description = "UPDATE_PROFILE_NOTES")
    @PutMapping("/profile/{id}")
    @ApiException(CREATE_MODEL_LIST_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:profile:update,empower:dataset:profile:update")
	public Result<Object> updateProfile(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, 
									  @PathVariable(value = "id") Integer profileId,
									  @RequestBody ModelProfile profile){
		return modelService.updateProfile(loginUser, profileId, profile);
	}
	
	@Operation(summary = "publish", description = "PUBLISH")
    @PostMapping("publish")
    @ApiException(PUBLISH_MODEL_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:publish,empower:dataset:publish")
	public Result<Object> publish(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, 
									@RequestBody PublishModelDto dto){
		return modelService.publish(loginUser, dto);
	}
	
	@Operation(summary = "unPublish", description = "UNPUBLISH")
    @PostMapping("/{id}/cancel")
    @ApiException(UNPUBLISH_MODEL_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:unPublish,empower:dataset:unPublish")
	public Result<Object> unPublish(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, 
								@PathVariable(value = "id") Integer modelId){
		return modelService.unPublish(loginUser, modelId);
	}
	
	@Operation(summary = "queryModelFileList", description = "QUERY_MODEL_FILE_LIST_NOTES")
    @GetMapping("/{id}/file")
    @ApiException(MODEL_QUERY_LIST_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:file:list,empower:dataset:file:list")
	public Result<Object> queryModelFileList(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, 
								@PathVariable(value = "id") Integer modelId,
								@RequestParam("pid") Integer pid){
		return Result.success(modelFileService.queryModelFileByModelId(modelId, pid));
	}
	
	@Operation(summary = "uploadFile", description = "UPLOADFILE")
    @PostMapping("/{id}/file/upload")
    @ApiException(UPLOAD_MODEL_FILE_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:file:upload,empower:dataset:file:upload")
	public Result<Object> uploadFile(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser, 
								@PathVariable(value = "id") Integer modelId,
								@RequestParam(value = "remark", required = false) String remark,
                                @RequestParam(value = "name") String name,
                                @RequestParam(value = "pid", required = false) Integer pid,
                                @RequestParam(value = "isDirectory", required = false) Integer isDirectory,
								@RequestParam(value = "file", required = false) MultipartFile file){
		ModelFileDto dto = new ModelFileDto();
		dto.setName(name);
		dto.setRemark(remark);
		if(isDirectory != null) {
			dto.setIsDirectory(isDirectory);
		}
		if(pid != null) {
			dto.setPid(pid);
		}
		return modelFileService.uploadFile(loginUser, modelId, dto, file);
	}
	
	@Operation(summary = "createFile", description = "CREATEFILE")
    @PostMapping("/{id}/file")
    @ApiException(CREATE_MODEL_FILE_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:file:create,empower:dataset:file:create")
	public Result<Object> createFile(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser,
								@PathVariable(value = "id") Integer modelId,
								@RequestBody ModelFileDto dto){
		return modelFileService.createFile(loginUser, modelId, dto);
	}
	
	@Operation(summary = "deleteFile", description = "DELETEFILE")
    @DeleteMapping("/{id}/file/{fileId}")
    @ApiException(DELETE_MODEL_FILE_ERROR)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    @PermissionCheck(permission = "empower:model:file:delete,empower:dataset:file:delete")
	public Result<Object> deleteFile(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser,
								@PathVariable(value = "id") Integer modelId,
								@PathVariable(value = "fileId") Integer fileId){
		return modelFileService.deleteFile(loginUser, modelId, fileId);
	}
	
	@Operation(summary = "downloadmodel", description = "DOWNLOAD_MODEL_INPUTSTREAM")
    @Parameters({
            @Parameter(name = "modelId", description = "MODEL_ID", required = true, schema = @Schema(implementation = Integer.class)),
    })
    @GetMapping(value = "/download")
    @ResponseStatus(HttpStatus.OK)
    @ApiException(RESOURCE_NOT_EXIST)
    @AccessLogAnnotation(ignoreRequestArgs = "loginUser")
    public Result<Object> downloadModel(@Parameter(hidden = true) @RequestAttribute(value = Constants.SESSION_USER) User loginUser,
                              @RequestParam(value = "modelId") Integer modelId, HttpServletResponse response) throws IOException {
		Result<Object> result = modelFileService.downloadFile(loginUser, modelId);
		if(result.isSuccess()) {
			String fileName = (String)result.getData();
			File file = new File(fileName);
	        BufferedInputStream inputStream = new BufferedInputStream(new FileInputStream(file));
	        response.setContentType("application/octet-stream");
	        response.setHeader("Content-Disposition","attachment;fileName=" + file.getName());
	        response.setHeader("Content-Length", String.valueOf(file.length()));
	        ServletOutputStream outputStream = response.getOutputStream();
	        IOUtils.copy(inputStream,outputStream);
	        return Result.success();
		}
		return result;
    }
}
